"""
可视化模块
负责所有图表和数据分析的可视化功能
"""

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import streamlit as st

from src.config import TOP_FEATURES_COUNT, SAMPLE_SIZE_FOR_VISUALIZATION


def plot_feature_importance(feature_importance_df, top_n=10):
    """
    绘制特征重要性图表
    
    Args:
        feature_importance_df (pd.DataFrame): 特征重要性数据框
        top_n (int): 显示前N个特征
    """
    top_features = feature_importance_df.head(top_n)
    
    fig = px.bar(
        top_features,
        x='Importance',
        y='Feature',
        orientation='h',
        color='Importance',
        title=f"Top {top_n} 影响房价的关键特征",
        labels={'Importance': '重要性得分', 'Feature': '特征名称'}
    )
    
    st.plotly_chart(fig, use_container_width=True)


def plot_feature_analysis(data, feature_name, target_name='SalePrice'):
    """
    绘制单个特征的详细分析图表
    
    Args:
        data (pd.DataFrame): 数据框
        feature_name (str): 特征名称
        target_name (str): 目标变量名称
    """
    # 采样数据以提高性能
    sample_data = data.sample(SAMPLE_SIZE_FOR_VISUALIZATION) if len(data) > SAMPLE_SIZE_FOR_VISUALIZATION else data
    
    # 创建子图
    fig = make_subplots(
        rows=2, cols=2,
        subplot_titles=(
            f'{feature_name}与房价散点图',
            f'{feature_name}分布直方图',
            f'{feature_name}箱线图',
            '房价分布'
        ),
        specs=[[{"secondary_y": False}, {"secondary_y": False}],
               [{"secondary_y": False}, {"secondary_y": False}]]
    )
    
    # 散点图
    fig.add_trace(
        go.Scatter(
            x=sample_data[feature_name],
            y=sample_data[target_name],
            mode='markers',
            marker=dict(color='blue', opacity=0.6),
            name=f'{feature_name} vs 房价'
        ),
        row=1, col=1
    )
    
    # 特征分布直方图
    fig.add_trace(
        go.Histogram(
            x=sample_data[feature_name],
            nbinsx=30,
            name=f'{feature_name}分布',
            marker_color='green'
        ),
        row=1, col=2
    )
    
    # 箱线图
    fig.add_trace(
        go.Box(
            y=sample_data[feature_name],
            name=f'{feature_name}箱线图',
            marker_color='orange'
        ),
        row=2, col=1
    )
    
    # 房价分布
    fig.add_trace(
        go.Histogram(
            x=sample_data[target_name],
            nbinsx=30,
            name='房价分布',
            marker_color='red'
        ),
        row=2, col=2
    )
    
    fig.update_layout(
        height=600,
        title_text=f"{feature_name}特征详细分析",
        showlegend=False
    )
    
    st.plotly_chart(fig, use_container_width=True)
    
    # 统计信息
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.metric("平均值", f"{sample_data[feature_name].mean():.2f}")
    with col2:
        st.metric("中位数", f"{sample_data[feature_name].median():.2f}")
    with col3:
        st.metric("标准差", f"{sample_data[feature_name].std():.2f}")
    with col4:
        st.metric("相关系数", f"{sample_data[feature_name].corr(sample_data[target_name]):.3f}")


def plot_model_comparison(model_performances):
    """
    绘制模型性能比较图表
    
    Args:
        model_performances (dict): 模型性能字典
    """
    # 创建比较图表
    fig_compare = make_subplots(
        rows=1, cols=2,
        subplot_titles=('R²分数比较', '均方误差比较'),
        specs=[[{"type": "bar"}, {"type": "bar"}]]
    )
    
    # R2分数比较
    fig_compare.add_trace(
        go.Bar(
            x=list(model_performances.keys()),
            y=[perf['r2'] for perf in model_performances.values()],
            name='R²分数',
            marker_color=['blue', 'red', 'green']
        ),
        row=1, col=1
    )
    
    # 均方误差比较
    fig_compare.add_trace(
        go.Bar(
            x=list(model_performances.keys()),
            y=[perf['mse'] for perf in model_performances.values()],
            name='均方误差',
            marker_color=['orange', 'purple', 'brown']
        ),
        row=1, col=2
    )
    
    fig_compare.update_layout(
        height=400,
        title_text="模型性能对比",
        showlegend=False
    )
    
    st.plotly_chart(fig_compare, use_container_width=True)
    
    # 性能指标表格
    performance_df = pd.DataFrame({
        '模型': list(model_performances.keys()),
        'R²分数': [perf['r2'] for perf in model_performances.values()],
        '均方误差': [perf['mse'] for perf in model_performances.values()],
        'RMSE': [perf['rmse'] for perf in model_performances.values()]
    })
    st.dataframe(performance_df)


def display_model_parameters(grid_searches):
    """
    显示模型的最佳参数配置
    
    Args:
        grid_searches (dict): 网格搜索对象字典
    """
    st.subheader("🔧 最佳参数配置")
    
    col1, col2, col3 = st.columns(3)
    
    model_names = {
        'rf': '随机森林',
        'svm': 'SVM', 
        'elastic': '弹性网络'
    }
    
    with col1:
        st.markdown("**随机森林最佳参数：**")
        for param, value in grid_searches['rf'].best_params_.items():
            st.write(f"- {param}: {value}")
    
    with col2:
        st.markdown("**SVM最佳参数：**")
        for param, value in grid_searches['svm'].best_params_.items():
            st.write(f"- {param}: {value}")
    
    with col3:
        st.markdown("**弹性网络最佳参数：**")
        for param, value in grid_searches['elastic'].best_params_.items():
            st.write(f"- {param}: {value}")


def plot_prediction_vs_actual(y_true, y_pred, model_name):
    """
    绘制预测值vs实际值的散点图
    
    Args:
        y_true (array): 实际值
        y_pred (array): 预测值
        model_name (str): 模型名称
    """
    fig = px.scatter(
        x=y_true,
        y=y_pred,
        title=f"{model_name} - 预测值 vs 实际值",
        labels={'x': '实际房价', 'y': '预测房价'}
    )
    
    # 添加对角线
    min_val = min(y_true.min(), y_pred.min())
    max_val = max(y_true.max(), y_pred.max())
    fig.add_trace(
        go.Scatter(
            x=[min_val, max_val],
            y=[min_val, max_val],
            mode='lines',
            name='完美预测线',
            line=dict(color='red', dash='dash')
        )
    )
    
    st.plotly_chart(fig, use_container_width=True)
